#include "Graph.hpp"

#include <map>
#include <set>
#include <stack>

template <typename T> auto createBipartiteGraph() {
  Graph<T> G(10);

  std::map<size_t, vector<std::pair<size_t, T>>> edges;
  edges[1] = {{2, 0}};
  edges[2] = {{1, 0}, {3, 0}, {8, 0}};
  edges[3] = {{2, 0}, {4, 0}};
  edges[4] = {{3, 0}, {6, 0}};
  edges[5] = {{7, 0}, {9, 0}};
  edges[6] = {{1, 0}, {4, 0}};
  edges[7] = {{5, 0}};
  edges[8] = {{2, 0}, {9, 0}};
  edges[9] = {{5, 0}};

  for (auto &i : edges)
    for (auto &j : i.second)
      G.addEdge(Edge<T>{i.first, j.first, j.second});

  return G;
}

template <typename T>
auto bipartiteCheck(const Graph<T>& G, size_t src) {
    std::stack<size_t> stack;
    std::set<size_t> visited;
    stack.push(src);
    enum class colors {NONE, RED, BLUE};
    colors currentColor {colors::BLUE};
    vector<colors> vertexColors(G.vertices(), colors::NONE);
    while (!stack.empty()) {
        auto currentVertex = stack.top();
        stack.pop();
        if (visited.find(currentVertex) == visited.end()) {
            visited.insert(currentVertex);
            vertexColors[currentVertex] = currentColor;
            if (currentColor == colors::RED) {
                cout << "Coloring vertex " << currentVertex
                     << " RED\n";
                currentColor = colors::BLUE;
            }
            else
            {
                cout << "Coloring vertex " << currentVertex
                     << " BLUE\n";
                currentColor = colors::RED;
            }
            for (const auto& e : G.edgesFromVertex(currentVertex))
                if (visited.find(e.dest) == visited.end())
                    stack.push(e.dest);
            //若已访问过，但是该顶点的颜色与将要设置的颜色不同
            //换言之是与父节点颜色相同，则表示该图为非二分图
        } else if ((vertexColors[currentVertex] == colors::BLUE &&
                    currentColor == colors::RED) ||
                   (vertexColors[currentVertex] == colors::RED &&
                    currentColor == colors::BLUE))
            return false;
    }
    return true;
}

template <typename T>
void testBipartite() {
    auto BG = createBipartiteGraph<T>();
    cout << BG << endl;
    auto ret = (bipartiteCheck<T>(BG, 1) ? "is " : "is not ");
    cout << "The graph " 
         << ret << "bipartite\n";
}

int main(int argc, const char** argv) {
    using T = uint32_t;
    testBipartite<T>();
    return 0;
}